
import streamlit as st
import asyncio
import nest_asyncio
from datetime import datetime

from agent.streaming_agent import StreamingChatbot, EventType, StreamingEvent

nest_asyncio.apply()

def render_event(event: StreamingEvent):

    timestamp = datetime.fromtimestamp(event.timestamp).strftime("%H:%M:%S.%f")[:-3]
    timestamp = f"<font color='red'> {timestamp} </font>"

    if event.type == EventType.THINKING_START:
        st.markdown(f"**{timestamp}** - 🤔 **开始思考...**", unsafe_allow_html=True)

    elif event.type == EventType.THINKING:
        thought = event.data.get("thought", "")
        st.markdown(f"```{thought}```")

    elif event.type == EventType.ACTION_PLANNED:
        action = event.data.get("action", "")
        action_input = event.data.get("input", "")
        st.markdown(f"**{timestamp}** - 📋 **动作规划：** `{action}`，**输入参数：** `{action_input}`", unsafe_allow_html=True)

    elif event.type == EventType.ACTION_START:
        action = event.data.get("action", "")
        st.markdown(f"**{timestamp}** - 🔧 **执行动作：** `{action}`", unsafe_allow_html=True)

    elif event.type == EventType.ACTION_RESULT:
        tool = event.data.get("tool", "")
        result = event.data.get("result", {})
        st.markdown(f"**{timestamp}** - 📤 **工具 {tool} 调用结果：**", unsafe_allow_html=True)
        if result.get("success"):
            data = result.get('data', 'No data')
            if isinstance(data, list) or isinstance(data, dict):
                st.json(data)
            else:
                st.markdown(str(data))
        else:
            st.error(f"Error: {result.get('error', 'Unknown error')}")

    elif event.type == EventType.OBSERVATION:
        observation = event.data.get("observation", "")
        st.markdown(f"**{timestamp}** - 👁️ **结果审查**", unsafe_allow_html=True)

    elif event.type == EventType.PLAN_CREATED:
        plan = event.data.get("plan", {})
        st.markdown(f"**{timestamp}** - 📋 **计划创建完成：**", unsafe_allow_html=True)
        st.markdown(f"**描述信息：** {plan.get('description', '无')}")

        steps = plan.get("steps", [])
        if steps:
            st.markdown("**步骤：**")
            for i, step in enumerate(steps, 1):
                st.markdown(f"{i}. **{step.get('tool', '未知')}**: {step.get('description', '无')}")

    elif event.type == EventType.PLAN_STEP_START:
        step_desc = event.data.get("step_description", "")
        tool = event.data.get("tool", "")
        plan_step = event.data.get("plan_step", 0)
        total_steps = event.data.get("total_steps", 0)
        st.markdown(f"**{timestamp}** - 🎯 **计划步骤 {plan_step}/{total_steps}：** `{tool}`", unsafe_allow_html=True)
        st.markdown(f"  _{step_desc}_")

    elif event.type == EventType.PLAN_STEP_COMPLETE:
        st.markdown(f"**{timestamp}** - ✅ **计划执行完毕**", unsafe_allow_html=True)

    elif event.type == EventType.REFLECTION_START:
        original_response = event.data.get("original_response", "")
        quality_threshold = event.data.get("quality_threshold", 0.7)
        st.markdown(f"**{timestamp}** - 🔍 **开始反思**", unsafe_allow_html=True)
        st.markdown(f"  **回答质量设定阈值：** {quality_threshold:.2f}")
        if original_response:
            st.markdown(
                f"  **当前回答预览：** {original_response[:100]}{'...' if len(original_response) > 100 else ''}")

    elif event.type == EventType.REFLECTION_CRITIQUE:
        iteration = event.data.get("iteration", 1)
        critique = event.data.get("critique", {})
        st.markdown(f"**{timestamp}** - 🤔 **反思批评-迭代轮次 {iteration}：**", unsafe_allow_html=True)

        quality = critique.get("overall_quality", 0.0)
        confidence = critique.get("confidence", 0.0)
        needs_refinement = critique.get("needs_refinement", False)

        if quality >= 0.8:
            st.success(f"回答质量评分：{quality:.2f} (优秀)")
        elif quality >= 0.7:
            st.info(f"回答质量评分：{quality:.2f} (良好)")
        elif quality >= 0.5:
            st.warning(f"回答质量评分：{quality:.2f} (需要改进)")
        else:
            st.error(f"回答质量评分：{quality:.2f} (低)")

        st.markdown(f"**置信度：** {confidence:.2f}， **需要完善：** {'是' if needs_refinement else '否'}")

        issues = critique.get("issues", [])
        if issues:
            st.markdown("**发现的问题：**")
            for i, issue in enumerate(issues, 1):
                issue_type = issue.get("type").title() if 'type' in issue else "未知"
                severity = issue.get("severity", "未知")
                description = issue.get("description", "无描述")
                suggestion = issue.get("suggestion", "无建议")

                cn_issue_type = {
                    "accuracy": "准确性",
                    "completeness": "完整性",
                    "clarity": "清晰度",
                    "relevance": "相关性",
                    "reasoning": "合理性",
                    "evidence": "充分性"
                }.get(str(issue_type).lower(), "未知")

                severity_emoji = {
                    "critical": "🔴",
                    "major": "🟠",
                    "minor": "🟡",
                    "suggestion": "🔵"
                }.get(severity.lower(), "⚪")

                ch_severity = {
                    "critical": "严重",
                    "major": "主要",
                    "minor": "次要",
                    "suggestion": "建议"
                }.get(severity.lower(), "次要")

                st.markdown(f"  {severity_emoji} **{cn_issue_type}** ({ch_severity}): {description}")
                st.markdown(f"  💡 **建议：** {suggestion}")

        strengths = critique.get("strengths", [])
        if strengths:
            st.markdown("**改进点：**")
            for strength in strengths:
                st.markdown(f"      ✅ {strength}")

        reasoning = critique.get("reasoning", "")
        if reasoning:
            st.markdown(f"**理由：** {reasoning}")

    elif event.type == EventType.REFLECTION_REFINEMENT:
        improvements = event.data.get("improvements", [])
        quality_improvement = event.data.get("quality_improvement", 0.0)
        st.markdown(f"**{timestamp}** - ✨ **优化响应：**", unsafe_allow_html=True)

        if quality_improvement > 0:
            st.success(f"回答质量得分改进：+{quality_improvement:.2f}")

        if improvements:
            st.markdown(f"**改进点({len(improvements)}个)：**")
            for improvement in improvements:
                st.markdown(f"      🔧 {improvement}")

    elif event.type == EventType.REFLECTION_COMPLETE:
        final_quality = event.data.get("final_quality_score", 0.0)
        threshold_met = event.data.get("threshold_met", False)
        st.markdown(f"**{timestamp}** - 🎉 **完成反思**", unsafe_allow_html=True)
        if threshold_met:
            st.success(f"✅ 达到回答质量阈值！最终得分：{final_quality:.2f}")
        else:
            st.warning(f"⚠️ 未达回答质量阈值！最终得分：{final_quality:.2f}")

st.set_page_config(
    page_title="通用智能体",
    page_icon="🧠",
    layout="wide",
    initial_sidebar_state="expanded"
)

if "chatbot" not in st.session_state:
    st.session_state.chatbot = StreamingChatbot(verbose=True, mode='plan_execute', enable_reflection=False)

async def streamlit_app():

    col1, col2 = st.columns([1, 1])
    mode = col1.selectbox(
        "智能体模式",
        ["plan_execute", "hybrid", "react"],
        index=0,
        label_visibility='collapsed',
        help="选择智能体处理模式"
    )

    enable_reflection = col2.checkbox(
        "启用反思",
        value=False,
        help="启用反思和响应优化"
    )

    if (mode != st.session_state.chatbot.agent.mode or
            enable_reflection != getattr(st.session_state.chatbot.agent, 'enable_reflection', True)
    ):
        st.session_state.chatbot = StreamingChatbot(verbose=False, mode=mode, enable_reflection=enable_reflection)
        st.rerun()

    try:
        with st.expander(label='智能体流程图', expanded=True):
            st.image(st.session_state.chatbot.agent.workflow_image())
    except:
        pass

    prompt = st.chat_input("可以问任何问题...")
    if prompt is None:
        return

    with st.chat_message("user"):
        st.write(prompt)

    final_response = None
    step = -1
    with st.chat_message("assistant"):
        response_placeholder = st.empty()
        with st.expander(label='智能体处理过程', expanded=True):
            with st.spinner("🤔 处理中 ...", show_time=True):
                async for event in st.session_state.chatbot.chat_stream(prompt):
                    if event.step != step:
                        st.markdown(f"#### 第 {event.step} 步", unsafe_allow_html=True)
                        step = event.step
                    render_event(event)
                    # if event.type == EventType.ERROR:
                    #     continue
                    await asyncio.sleep(0.01)
                    if event.type == EventType.COMPLETE:
                        final_response = event.data.get("response")['output']

    if final_response is None:
        final_response = f"抱歉, 遇到错误：{event.data.get('error', '未知错误')}"
    response_placeholder.write(final_response)


if __name__ == "__main__":
    asyncio.run(streamlit_app())